package com.by.rpc.pool;

import com.by.log.Logger;
import com.by.log.LoggerFactory;
import com.by.rpc.netty.RpcClient;

import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * RPC连接池，用于管理和复用RPC客户端连接
 */
public class RpcConnectionPool {
    private static final Logger logger = LoggerFactory.getLogger(RpcConnectionPool.class);

    private final int maxConnections;
    private final int minConnections;
    private final long maxIdleTimeMs;
    private final String host;
    private final int port;

    private final BlockingQueue<RpcClient> connectionQueue;
    private final AtomicInteger activeConnections = new AtomicInteger(0);
    private volatile boolean isClosed = false;

    // 独立的线程池，用于执行阻塞的连接获取和初始化操作
    private final ExecutorService connectionExecutor;

    /**
     * 构造函数
     * @param host 服务端主机
     * @param port 服务端端口
     * @param maxConnections 最大连接数
     * @param minConnections 最小连接数
     * @param maxIdleTimeMs 最大空闲时间(毫秒)
     */
    public RpcConnectionPool(String host, int port, int maxConnections, int minConnections, long maxIdleTimeMs) {
        this.host = host;
        this.port = port;
        this.maxConnections = maxConnections;
        this.minConnections = minConnections;
        this.maxIdleTimeMs = maxIdleTimeMs;
        this.connectionQueue = new LinkedBlockingQueue<>(maxConnections);

        if (minConnections > maxConnections) {
            throw new IllegalArgumentException("minConnections cannot be greater than maxConnections");
        }

        this.connectionExecutor = new ThreadPoolExecutor(
                1, maxConnections, 60L, TimeUnit.SECONDS,
                new LinkedBlockingQueue<>(),
                new ThreadFactory() {
                    private final AtomicInteger counter = new AtomicInteger(0);
                    @Override
                    public Thread newThread(Runnable r) {
                        return new Thread(r, String.format("RpcConnectionPool-%s:%d-%d", host, port, counter.incrementAndGet()));
                    }
                });

        logger.info("Creating RpcConnectionPool for {}:{} with maxConnections={}, minConnections={}, maxIdleTimeMs={}",
                host, port, maxConnections, minConnections, maxIdleTimeMs);

        initializeMinConnections();
    }

    /**
     * 初始化最小连接数（异步执行，不阻塞主线程）
     */
    private void initializeMinConnections() {
        connectionExecutor.submit(() -> {
            logger.debug("Initializing minimum connections: {}", minConnections);
            for (int i = 0; i < minConnections; i++) {
                try {
                    RpcClient client = new RpcClient(host, port); // 构造函数不再阻塞
                    client.connectAsync().get(maxIdleTimeMs, TimeUnit.MILLISECONDS); // 阻塞等待异步连接完成
                    if (isConnectionActive(client)) {
                        connectionQueue.offer(client);
                        logger.debug("Initialized connection {}/{}", i + 1, minConnections);
                    } else {
                        logger.warn("Failed to initialize connection {}/{}", i + 1, minConnections);
                        client.close(); // 确保关闭失败的客户端
                    }
                } catch (Exception e) {
                    logger.error("Error initializing connection {}/{}", i + 1, minConnections, e);
                }
            }
            logger.debug("Initialized {} connections", connectionQueue.size());
        });
    }

    /**
     * 同步获取连接（此方法仍然是同步阻塞的，但它会等待异步连接完成）
     */
    public RpcClient getConnection() throws InterruptedException {
        if (isClosed) {
            throw new IllegalStateException("Connection pool is closed");
        }

        RpcClient client = connectionQueue.poll();

        if (client == null) {
            int currentActive = activeConnections.get();
            if (currentActive < maxConnections) {
                logger.debug("Creating new connection. Current active: {}", currentActive);
                client = new RpcClient(host, port);
                try {
                    client.connectAsync().get(5, TimeUnit.SECONDS); // 阻塞等待异步连接完成
                } catch (Exception e) {
                    logger.error("Failed to establish new connection to {}:{}", host, port, e);
                    client.close(); // 连接失败，关闭客户端
                    throw new RuntimeException("Failed to establish active connection to " + host + ":" + port, e);
                }
            } else {
                logger.debug("Max connections reached, waiting for available connection");
                client = connectionQueue.poll(5, TimeUnit.SECONDS); // 等待从队列获取

                if (client == null) {
                    logger.debug("Timeout waiting for connection, creating new one");
                    client = new RpcClient(host, port);
                    try {
                        client.connectAsync().get(5, TimeUnit.SECONDS); // 阻塞等待异步连接完成
                    } catch (Exception e) {
                        logger.error("Failed to establish new connection to {}:{}", host, port, e);
                        client.close();
                        throw new RuntimeException("Failed to establish active connection to " + host + ":" + port, e);
                    }
                }
            }
        } else {
            if (!isConnectionActive(client)) {
                logger.warn("Found inactive client in pool, creating new one");
                try {
                    client.close();
                } catch (Exception e) {
                    logger.warn("Error closing inactive client", e);
                }
                client = new RpcClient(host, port);
                try {
                    client.connectAsync().get(5, TimeUnit.SECONDS); // 阻塞等待异步连接完成
                } catch (Exception e) {
                    logger.error("Failed to establish new connection to {}:{}", host, port, e);
                    client.close();
                    throw new RuntimeException("Failed to establish active connection to " + host + ":" + port, e);
                }
            }
        }

        if (client == null || !isConnectionActive(client)) {
            // 理论上到这里，client应该已经是活跃的，除非上面捕获异常后没有正确处理
            throw new RuntimeException("Failed to establish active connection to " + host + ":" + port);
        }

        activeConnections.incrementAndGet();
        logger.debug("Got connection. Active connections: {}", activeConnections.get());
        return client;
    }

    /**
     * 异步获取连接，供外部调用
     */
    public CompletableFuture<RpcClient> getConnectionAsync() {
        logger.debug("Async getting connection from pool.");
        return CompletableFuture.supplyAsync(() -> {
            try {
                return getConnection(); // 仍在 connectionExecutor 线程中执行阻塞，但对外部调用者是异步的
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                throw new CompletionException("Failed to get a connection due to interruption", e);
            }
        }, connectionExecutor); // 确保在 connectionExecutor 线程池中执行
    }

    private boolean isConnectionActive(RpcClient client) {
        if (client == null) {
            return false;
        }
        boolean active = client.isActive();
        logger.debug("Checking connection status for {}:{} - active: {}", host, port, active);
        return active;
    }

    public void returnConnection(RpcClient client) {
        logger.debug("Returning connection to pool. Queue size: {}, Active connections: {}",
                connectionQueue.size(), activeConnections.get());

        if (isClosed) {
            if (client != null) {
                try {
                    client.close();
                } catch (Exception e) {
                    logger.warn("Error closing client while pool is closed", e);
                }
            }
            return;
        }

        if (client != null && isConnectionActive(client)) {
            // 确保只有活跃且未超时的连接才被返回到队列
            // 可以在此处添加一个逻辑，检查客户端是否长时间未使用（空闲超时），如果超时则关闭而不返回
            if (System.currentTimeMillis() - client.getLastUseTime() > maxIdleTimeMs) { // 假设 RpcClient 有 getLastUseTime() 方法
                logger.debug("Client has been idle for too long, closing it instead of returning to pool.");
                try {
                    client.close();
                } catch (Exception e) {
                    logger.warn("Error closing idle client", e);
                }
            } else if (connectionQueue.size() < maxConnections) {
                connectionQueue.offer(client);
                logger.debug("Returned connection to queue");
            } else {
                logger.debug("Connection queue is full, closing connection");
                try {
                    client.close();
                } catch (Exception e) {
                    logger.warn("Error closing client when queue is full", e);
                }
            }
        } else {
            logger.debug("Client is inactive or null, closing and decrementing active count");
            if (client != null) {
                try {
                    client.close();
                } catch (Exception e) {
                    logger.warn("Error closing inactive client", e);
                }
            }
        }

        activeConnections.decrementAndGet();
        logger.debug("Connection returned. Active connections: {}", activeConnections.get());
    }

    public void close() {
        logger.info("Closing RpcConnectionPool for {}:{}", host, port);
        isClosed = true;

        RpcClient client;
        while ((client = connectionQueue.poll()) != null) {
            try {
                client.close();
            } catch (Exception e) {
                logger.warn("Error closing client during pool shutdown", e);
            }
        }

        connectionExecutor.shutdown();
        try {
            if (!connectionExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
                connectionExecutor.shutdownNow();
            }
        } catch (InterruptedException e) {
            connectionExecutor.shutdownNow();
            Thread.currentThread().interrupt();
        }

        logger.info("RpcConnectionPool for {}:{} closed", host, port);
    }
}